import pandas as pd
import matplotlib.pylab as plt
from scipy.stats import chi2_contingency


model_set = pd.read_csv('model_set.csv')
test_set = pd.read_csv('test_set.csv')

def qcut(data, col, num_bins=20, bins=None):
    '''
    数据分箱
    input:
        data, 数据集
        col, 需要分箱的列
        num_bins, 等频分箱个数，默认20
        bins, 手动传入分箱
    return:
        
    '''
    #如果没有传入分箱边界，则等频分箱
    if not bins:
        data['qcut'], low_bounday = pd.qcut(data[col], num_bins, retbins=True, duplicates='drop')
    else:
        data['qcut'], low_bounday = pd.cut(data[col], bins, retbins=True)
    #去除分箱各个分类的个数
    
    grouped = data.groupby('qcut').SeriousDlqin2yrs.agg(['sum', 'count'])
    grouped.columns = ['count_1', 'total']
    grouped['count_0'] = grouped['total'] - grouped['count_1']
    
    bins_array = [*zip(low_bounday, low_bounday[1:], grouped.count_0, grouped.count_1)]

    return bins_array

def cal_woe_iv(bins_array):
    '''
    计算每个分箱的woe值
    input:
        bins_array: 上一步计算好的每个分箱包含 count_0, count_1
    '''
    col = ['low', 'up', 'count_0', 'count_1']
    df_bins = pd.DataFrame(bins_array, columns = col)
    df_bins['total'] = df_bins['count_0'] + df_bins['count_1']
    df_bins['percentage'] = df_bins['total'] / df_bins['total'].sum()
    #分箱内坏的比例
    df_bins['bad_rate'] = df_bins.count_1 / df_bins.total
    #坏的占总样本
    df_bins['bad%'] = df_bins.count_1 / df_bins.count_1.sum()
    df_bins['good%'] = df_bins.count_0 / df_bins.count_0.sum()
    df_bins['woe'] = np.log(df_bins['good%'] / df_bins['bad%'])
    #计算iv值
    df_bins['good-bad'] = df_bins['good%'] - df_bins['bad%']
    iv = (df_bins['good-bad'] * df_bins['woe']).sum()
    return df_bins, iv

def plot_for_bins(data, col, n=2):
    
    bins_array = qcut(data, col)
    IV = []
    x_axis = []
    #当分箱数量大于指定分箱时，合并p值最大的分箱
    while len(bins_array) > n:
        p_value_list = []
        for i in range(len(bins_array) - 1):
            #取count_0, count_1的数量进行卡方检验
            x1 = bins_array[i][-2:]
            x2 = bins_array[i][-2:]
            p_value = chi2_contingency([x1,x2])[1]
            p_value_list.append(p_value)
        
        i = np.argmax(p_value_list)
        #合并分箱
        bins_array[i:i+2] = [(
                bins_array[i][0],#分箱的下界
                bins_array[i+1][1],#下一个分箱的上界
                bins_array[i][2] + bins_array[i+1][2],
                bins_array[i][3] + bins_array[i+1][3])]
        
        bins_df, iv = cal_woe_iv(bins_array)
        IV.append(iv)
        x_axis.append(len(bins_array))
        
    plt.figure()
    plt.plot(x_axis, IV)
    plt.xticks(x_axis)
    plt.title(col)
    plt.show()


###画图确定分箱
for i in model_set.columns[1:]:
    try:
        plot_for_bins(model_set, i)
    except:
        print (i)